#Author: Tim Henderson
#Maintainer: Tim Henderson @ 
#   tim.tadh@gmail.com
#   400 Braemer Ct. Gahanna OH 43230 USA
#Purpose: To allow the user to save files in an encrypted archive
#Copyright: (c)2006 Tim Henderson. All Rights Reserved. This module is part of 
#   the program secureData.
#License: secureData and all of its component parts are licensed under the 
#   terms and conditions of the GNU General Public License Version 2. Later versions
#   of this license do not apply to secureData or its component parts. If you have 
#   not recieved a copy of the license with your software please contact the Free 
#   Software Foundation at:
#       51 Franklin St, Fifth Floor, Boston, MA 02110-1301 USA
#   or you may also locate the license at the following internet address
#       http://www.gnu.org/licenses/gpl.txt

from Crypto.Cipher import AES
from Crypto.Hash import MD2, MD5
import os, stat
import reDDB, nDDB
import cStringIO
import glob

class EFAerror(Exception): pass
class ControlIndexTooLongError(EFAerror): pass
class FileAlreadyInArchiveError(EFAerror): pass
class FileDoesntExistError(EFAerror): pass
class FileNotInArchiveError(EFAerror): pass
class BadControlDDB(EFAerror): pass
class InvalidArchiveError(EFAerror): pass

def hMD2gen(s):
    md2 = MD2.new()
    md2.update(s)
    md2.update(md2.hexdigest())
    c = md2.hexdigest()
    while True:
        yield c
        md2.update(c)
        c = md2.hexdigest()
        
class EncryptedFA(object):
    
    CLM = 12
    BLOCK = 2**20
    
    def __init__(self, path):
        self.path = path
        if self._pathExists(path):
            self._readIn()
        else:
            self._new()
    
    def _pathExists(self, path):
        if len(glob.glob(path)) > 0:
            return True;
        else:
            return False;
    
    def _readIn(self):
        fmrb = open(self.path, 'rb')
        try:
            self.controlLen = self._readControlLen(fmrb)
        except:
            fmrb.close()
            self.controlLen = None
            self.control = None
            raise InvalidArchiveError, "The CLM mark was unreadable"
        self.control = self._readControlDDB(self.controlLen, fmrb)
        fmrb.close()
    
    def _readControlLen(self, f):
        mylen = int(os.stat(self.path)[stat.ST_SIZE])
        f.seek(mylen-self.CLM)
        slen = f.read(self.CLM)
        return int(slen)
    
    def _readControlDDB(self, clen, f):
        mylen = int(os.stat(self.path)[stat.ST_SIZE])
        f.seek(mylen - self.CLM - clen)
        ddb = f.read(clen)
        try:
            d = nDDB.decodeDDB(ddb)
            if d == None:
                raise BadControlDDB, 'The control index was corrupted'
            return d
        except:
            raise BadControlDDB, 'The control index was corrupted'
            return None
    
    def readFile(self, oF, name, keyGen, added):
        iF = open(self.path, 'rb')
        if self.control.has_key(name):
            flen = int(self.control[name]['len'])
            if flen < self.BLOCK:
                
                aes = AES.new(keyGen.next(), AES.MODE_ECB)
                spot = aes.decrypt(self._readSpot(int(self.control[name]['start']),
                                        flen, iF))
                if added != 0:
                    oF.write(spot[:-added])
                else:
                    oF.write(spot)
                del aes
            else:
                clen = flen
                tread = 0
                while clen - self.BLOCK > 0:
                    aes = AES.new(keyGen.next(), AES.MODE_ECB)
                    oF.write(aes.decrypt(self._readSpot(int(self.control[name]['start']) + tread,
                                        self.BLOCK, iF)))
                    del aes
                    tread += self.BLOCK
                    clen -= self.BLOCK
                aes = AES.new(keyGen.next(), AES.MODE_ECB)
                spot = aes.decrypt(self._readSpot(int(self.control[name]['start']) + tread,
                                        clen, iF))
                if added != 0:
                    oF.write(spot[:-added])
                else:
                    oF.write(spot)
                del aes
        else:
            raise FileNotInArchiveError, "The file %s is not in the archive" % name
    
    def _readSpot(self, start, slen, f):
        f.seek(start)
        spot = f.read(slen)
        return spot
    
    def _resetControlLen(self):
        self.controlLen = len(nDDB.makeAdvanceDDB(self.control))
    
    def _new(self):
        fmwb = open(self.path, 'wb')
        self.control = {}
        self.controlLen = len(nDDB.makeAdvanceDDB(self.control))
        fmwb.close()
        self._writeControl()
    
    def __len__(self):
        return self.__customControlLen(self.controlLen)
    
    def __customControlLen(self, clen):
        length = self.CLM
        length += self.controlLen
        for name in self.control:
            length += int(self.control[name]['len'])
        return length
    
    def _shiftLeft(self, fmrbp, start, amt, tlen):
        wcur = start
        rcur = wcur + amt
        fmrbp.seek(rcur)
        ch = fmrbp.read(self.BLOCK)
        while len(ch) == self.BLOCK:
            fmrbp.seek(wcur)
            fmrbp.write(ch)
            fmrbp.flush()
            wcur += self.BLOCK
            rcur += self.BLOCK
            fmrbp.seek(rcur)
            ch = fmrbp.read(self.BLOCK)
        fmrbp.seek(wcur)
        fmrbp.write(ch)
        fmrbp.flush()
        
        if tlen != None: fmrbp.truncate(tlen)
    
    def _shiftRight(self, fmrbp, start, amt, tlen):
        if tlen - start < self.BLOCK: block = tlen - start
        else: block = self.BLOCK
        rcur = tlen
        wcur = rcur + amt
        ch = ''
        while rcur - block > start:
            rcur -= block
            wcur -= block
            
            fmrbp.seek(rcur)
            ch = fmrbp.read(block)
            fmrbp.seek(wcur)
            fmrbp.write(ch)
            fmrbp.flush()
            
        a = rcur - start
        rcur = start
        wcur = start + amt
        
        fmrbp.seek(rcur)
        ch = fmrbp.read(a)
        fmrbp.seek(wcur)
        fmrbp.write(ch)
        fmrbp.flush()
    
    def _deleteFileFromControl(self, fname):
        flen = int(self.control[fname]['len'])
        fstart = int(self.control[fname]['start'])
        
        olen = len(self)
        olenslen = len(str(olen))
        oclen = self.controlLen
        mlen = len(self.control[fname]['start'])
        
        del self.control[fname]
        
        for key in self.control:
            cs = int(self.control[key]['start'])
            if cs > fstart:
                self.control[key]['start'] = str(cs - flen)
        
        self._resetControlLen()
    
    def deleteFile(self, fname):
        if self.control.has_key(fname):
            flen = int(self.control[fname]['len'])
            fstart = int(self.control[fname]['start'])
            
            oclen = self.controlLen
            olen = len(self)
            self._deleteFileFromControl(fname)
            
            fmrbp = open(self.path, 'rb+')
            #print fstart, flen, len(self)
            
            self._shiftLeft(fmrbp, fstart, flen, (olen - oclen - self.CLM))
            self._writeControl(fmrbp)
            
            fmrbp.close()
        else:
            raise FileNotInArchiveError, "The file %s is not in the archive" % fname
    
    def _addFileToControl(self, fname, flen):
        # flen = int(os.stat(abspath)[stat.ST_SIZE])
        # if flen%16 !=0:
            # flen += (16 - flen%16)
        # fname = os.path.basename(abspath)
        
        olen = len(self)
        olenslen = len(str(olen)) + 2
        oclen = self.controlLen + self.CLM
        
        if self.control.has_key(fname): 
            raise FileAlreadyInArchiveError, 'You already have a file of the name %s in the EFA' % (fname,)
        
        self.control.update({fname:{'len':flen, 'start':str(olen - oclen)}})
        
        self._resetControlLen()
    
    def addFile(self, iF, fname, flen, keyGen):
        #this needs to be changed to reading in by blocks bwe
        # path = os.path.abspath(path)
        # fname = os.path.basename(path)
        # 
        # if not self._pathExists(path):
            # raise FileDoesntExistError, 'The path %s does not exist' % (path,)
        
        #flen = int(os.stat(path)[stat.ST_SIZE])
        if flen%16 !=0:
            flen += (16 - flen%16)
        oldLen = self.CLM + self.controlLen
        tOldLen = len(self)
        self._addFileToControl(fname, flen)
        
        fmrbp = open(self.path, 'rb+')
        #print flen
        
        fmrbp.seek(tOldLen - oldLen)        
        
        #iF = open(path, 'rb')
        if flen < self.BLOCK:
            aes = AES.new(keyGen.next(), AES.MODE_ECB)
            spot, added = self._appendSpaces(iF.read())
            fmrbp.write(aes.encrypt(spot))
            del aes
            fmrbp.flush()
        else:
            clen = flen
            while clen - self.BLOCK > 0:
                aes = AES.new(keyGen.next(), AES.MODE_ECB)
                fmrbp.write(aes.encrypt(iF.read(self.BLOCK)))
                del aes
                fmrbp.flush()
                clen -= self.BLOCK
            aes = AES.new(keyGen.next(), AES.MODE_ECB)
            spot, added = self._appendSpaces(iF.read(clen))
            fmrbp.write(aes.encrypt(spot))
            del aes
            fmrbp.flush()

        self._writeControl(fmrbp)
        fmrbp.close()
        
        return added
    
    def _appendSpaces(self, plaintext):
        x = 0
        if len(plaintext)%16 != 0:
            x = 16 - len(plaintext)%16
            plaintext += ' '*x
        return plaintext, x
    
    def _writeControl(self, fmrbp=None):
        if fmrbp == None: fmrbp = open(self.path, 'rb+')
        cddb = nDDB.makeAdvanceDDB(self.control)
        cddbLen = str(len(cddb))
        
        if len(cddbLen) > self.CLM: 
            raise ControlIndexTooLongError, 'The control index has exceeded the allotted space of %s bytes' % ((self.CLM * '9'),)
        
        cddbLen = (self.CLM - len(cddbLen))*'0' + cddbLen
        
        mylen = len(self) - self.CLM - self.controlLen
        fmrbp.seek(mylen)
        fmrbp.write(cddb)
        fmrbp.flush()
        fmrbp.seek(mylen + len(cddb))
        fmrbp.write(cddbLen)
        fmrbp.flush()
        
        if fmrbp == None: fmrbp.close()